# -*- coding: utf-8 -*-
"""
股票筛选模块 - 根据指定条件筛选股票
"""
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
import logging
from typing import List, Dict, Tuple
from data_fetcher import get_data_fetcher
from config import SCREENING_CRITERIA, SCORING_WEIGHTS
import time

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class StockScreener:
    """股票筛选器"""
    
    def __init__(self):
        self.data_fetcher = get_data_fetcher()
        self.criteria = SCREENING_CRITERIA
        self.weights = SCORING_WEIGHTS
    
    def screen_stocks(self) -> pd.DataFrame:
        """根据条件筛选股票"""
        logger.info("开始筛选股票...")
        
        # 1. 获取股票列表
        stock_list = self.data_fetcher.get_stock_list()
        if stock_list.empty:
            logger.error("无法获取股票列表")
            return pd.DataFrame()
        
        # 2. 获取实时数据
        realtime_data = self.data_fetcher.get_stock_realtime_data([])
        if realtime_data.empty:
            logger.error("无法获取实时数据")
            return pd.DataFrame()
        
        # 3. 合并数据
        merged_data = pd.merge(
            stock_list, 
            realtime_data, 
            left_on='code', 
            right_on='代码', 
            how='inner'
        )
        
        logger.info(f"合并后数据量: {len(merged_data)}")
        
        # 4. 应用筛选条件
        filtered_stocks = self._apply_screening_criteria(merged_data)
        
        # 5. 获取详细信息并评分
        scored_stocks = self._score_stocks(filtered_stocks)
        
        # 6. 排序并选出前N只
        top_stocks = scored_stocks.nlargest(self.criteria['top_stocks'] * 3, 'total_score')
        
        logger.info(f"筛选完成，共找到 {len(top_stocks)} 只潜力股票")
        
        return top_stocks
    
    def _apply_screening_criteria(self, data: pd.DataFrame) -> pd.DataFrame:
        """应用筛选条件"""
        logger.info("应用筛选条件...")
        
        original_count = len(data)
        
        # 筛选条件
        conditions = []
        
        # 1. 市值筛选（总市值 >= 50亿）
        if '总市值' in data.columns:
            conditions.append(data['总市值'] >= self.criteria['market_cap_min'] * 100000000)
            logger.info(f"市值筛选: >= {self.criteria['market_cap_min']}亿")
        
        # 2. 股价筛选（4-20元）
        if '最新价' in data.columns:
            conditions.append(
                (data['最新价'] >= self.criteria['price_min']) & 
                (data['最新价'] <= self.criteria['price_max'])
            )
            logger.info(f"股价筛选: {self.criteria['price_min']}-{self.criteria['price_max']}元")
        
        # 3. 市盈率筛选（市盈率 > 0，即盈利股票）
        if '市盈率-动态' in data.columns:
            pe_condition = (data['市盈率-动态'] > self.criteria['pe_ratio_min']) & (data['市盈率-动态'] < 100)
            conditions.append(pe_condition)
            logger.info("市盈率筛选: 盈利股票且PE < 100")
        
        # 4. 排除ST股票
        if 'name' in data.columns:
            conditions.append(~data['name'].str.contains('ST|退|\\*', na=False))
            logger.info("排除ST和退市股票")
        
        # 应用所有条件
        if conditions:
            final_condition = conditions[0]
            for condition in conditions[1:]:
                final_condition = final_condition & condition
            
            filtered_data = data[final_condition].copy()
        else:
            filtered_data = data.copy()
        
        logger.info(f"筛选结果: {original_count} -> {len(filtered_data)} 只股票")
        
        return filtered_data
    
    def _score_stocks(self, stocks: pd.DataFrame) -> pd.DataFrame:
        """为股票评分"""
        logger.info("开始为股票评分...")
        
        scored_stocks = stocks.copy()
        
        # 初始化评分字段
        scored_stocks['technical_score'] = 0
        scored_stocks['fundamental_score'] = 0
        scored_stocks['momentum_score'] = 0
        scored_stocks['dividend_score'] = 0
        scored_stocks['total_score'] = 0
        
        for idx, row in scored_stocks.iterrows():
            stock_code = row['code']
            
            try:
                # 1. 技术指标评分
                tech_score = self._calculate_technical_score(stock_code)
                scored_stocks.at[idx, 'technical_score'] = tech_score
                
                # 2. 基本面评分
                fund_score = self._calculate_fundamental_score(row)
                scored_stocks.at[idx, 'fundamental_score'] = fund_score
                
                # 3. 动量评分
                momentum_score = self._calculate_momentum_score(stock_code)
                scored_stocks.at[idx, 'momentum_score'] = momentum_score
                
                # 4. 分红评分
                dividend_score = self._calculate_dividend_score(stock_code)
                scored_stocks.at[idx, 'dividend_score'] = dividend_score
                
                # 5. 总评分
                total_score = (
                    tech_score * self.weights['technical_score'] +
                    fund_score * self.weights['fundamental_score'] +
                    momentum_score * self.weights['momentum_score'] +
                    dividend_score * self.weights['dividend_score']
                )
                scored_stocks.at[idx, 'total_score'] = total_score
                
                logger.info(f"股票 {stock_code} 评分完成: {total_score:.2f}")
                
                # 避免请求过于频繁
                time.sleep(0.5)
                
            except Exception as e:
                logger.warning(f"股票 {stock_code} 评分失败: {e}")
                continue
        
        return scored_stocks
    
    def _calculate_technical_score(self, stock_code: str) -> float:
        """计算技术指标评分"""
        try:
            # 获取历史数据
            hist_data = self.data_fetcher.get_stock_historical_data(
                stock_code, 
                start_date=(datetime.now() - timedelta(days=180)).strftime("%Y%m%d")
            )
            
            if hist_data.empty or len(hist_data) < 60:
                return 0
            
            # 计算技术指标
            score = 0
            
            # 1. 均线趋势（30%）
            ma5 = hist_data['close'].rolling(5).mean()
            ma20 = hist_data['close'].rolling(20).mean()
            ma60 = hist_data['close'].rolling(60).mean()
            
            current_price = hist_data['close'].iloc[-1]
            
            # 价格相对均线位置
            if current_price > ma5.iloc[-1] > ma20.iloc[-1] > ma60.iloc[-1]:
                score += 30  # 多头排列
            elif current_price > ma20.iloc[-1]:
                score += 20  # 中期向上
            elif current_price > ma60.iloc[-1]:
                score += 10  # 长期向上
            
            # 2. 成交量趋势（20%）
            vol_ma5 = hist_data['volume'].rolling(5).mean()
            vol_ma20 = hist_data['volume'].rolling(20).mean()
            
            if vol_ma5.iloc[-1] > vol_ma20.iloc[-1] * 1.2:
                score += 20  # 成交量放大
            elif vol_ma5.iloc[-1] > vol_ma20.iloc[-1]:
                score += 10
            
            # 3. 相对强弱（25%）
            returns = hist_data['close'].pct_change(20).iloc[-1]
            if returns > 0.1:
                score += 25  # 强势上涨
            elif returns > 0.05:
                score += 15
            elif returns > 0:
                score += 10
            
            # 4. 波动率（25%）
            volatility = hist_data['close'].pct_change().rolling(20).std().iloc[-1]
            if 0.02 < volatility < 0.05:  # 适中波动率
                score += 25
            elif volatility < 0.02:
                score += 15  # 低波动
            
            return min(score, 100)
            
        except Exception as e:
            logger.warning(f"技术指标评分失败 {stock_code}: {e}")
            return 0
    
    def _calculate_fundamental_score(self, stock_data: pd.Series) -> float:
        """计算基本面评分"""
        try:
            score = 0
            
            # 1. 市盈率评分（30%）
            if '市盈率-动态' in stock_data:
                pe = stock_data['市盈率-动态']
                if pd.notna(pe) and pe > 0:
                    if 5 <= pe <= 20:
                        score += 30  # 合理PE
                    elif 20 < pe <= 30:
                        score += 20
                    elif pe <= 5:
                        score += 10
            
            # 2. 市净率评分（20%）
            if '市净率' in stock_data:
                pb = stock_data['市净率']
                if pd.notna(pb) and pb > 0:
                    if 1 <= pb <= 3:
                        score += 20  # 合理PB
                    elif 3 < pb <= 5:
                        score += 15
                    elif pb < 1:
                        score += 10
            
            # 3. 市值评分（25%）
            if '总市值' in stock_data:
                market_cap = stock_data['总市值'] / 100000000  # 转换为亿元
                if 50 <= market_cap <= 200:
                    score += 25  # 中等市值
                elif 200 < market_cap <= 500:
                    score += 20
                elif market_cap > 500:
                    score += 15
            
            # 4. 换手率评分（25%）
            if '换手率' in stock_data:
                turnover = stock_data['换手率']
                if pd.notna(turnover):
                    if 1 <= turnover <= 5:
                        score += 25  # 适度活跃
                    elif 5 < turnover <= 10:
                        score += 20
                    elif turnover < 1:
                        score += 10
            
            return min(score, 100)
            
        except Exception as e:
            logger.warning(f"基本面评分失败: {e}")
            return 0
    
    def _calculate_momentum_score(self, stock_code: str) -> float:
        """计算动量评分"""
        try:
            # 获取历史数据
            hist_data = self.data_fetcher.get_stock_historical_data(
                stock_code,
                start_date=(datetime.now() - timedelta(days=120)).strftime("%Y%m%d")
            )
            
            if hist_data.empty or len(hist_data) < 30:
                return 0
            
            score = 0
            
            # 1. 短期动量（40%）
            returns_5d = (hist_data['close'].iloc[-1] / hist_data['close'].iloc[-6] - 1) * 100
            if returns_5d > 5:
                score += 40
            elif returns_5d > 2:
                score += 30
            elif returns_5d > 0:
                score += 20
            
            # 2. 中期动量（35%）
            returns_20d = (hist_data['close'].iloc[-1] / hist_data['close'].iloc[-21] - 1) * 100
            if returns_20d > 10:
                score += 35
            elif returns_20d > 5:
                score += 25
            elif returns_20d > 0:
                score += 15
            
            # 3. 相对大盘强度（25%）
            # 这里简化处理，实际应该对比大盘指数
            if returns_20d > 3:  # 假设大盘平均涨幅3%
                score += 25
            elif returns_20d > 0:
                score += 15
            
            return min(score, 100)
            
        except Exception as e:
            logger.warning(f"动量评分失败 {stock_code}: {e}")
            return 0
    
    def _calculate_dividend_score(self, stock_code: str) -> float:
        """计算分红评分"""
        try:
            dividend_data = self.data_fetcher.get_dividend_info(stock_code)
            
            if dividend_data.empty:
                return 0 if self.criteria['dividend_required'] else 30  # 如果不要求分红，给基础分
            
            score = 0
            
            # 1. 分红连续性（50%）
            recent_years = 3
            recent_dividends = dividend_data.head(recent_years)
            
            if len(recent_dividends) >= recent_years:
                score += 50  # 连续分红
            elif len(recent_dividends) >= 2:
                score += 35
            elif len(recent_dividends) >= 1:
                score += 20
            
            # 2. 分红率（50%）
            if not dividend_data.empty:
                try:
                    # 计算股息率（简化）
                    latest_dividend = dividend_data.iloc[0]
                    if '股息率' in dividend_data.columns:
                        dividend_yield = latest_dividend['股息率']
                        if dividend_yield > 3:
                            score += 50
                        elif dividend_yield > 2:
                            score += 35
                        elif dividend_yield > 1:
                            score += 20
                except:
                    pass
            
            return min(score, 100)
            
        except Exception as e:
            logger.warning(f"分红评分失败 {stock_code}: {e}")
            return 0 if self.criteria['dividend_required'] else 30

if __name__ == "__main__":
    # 测试股票筛选
    screener = StockScreener()
    top_stocks = screener.screen_stocks()
    
    if not top_stocks.empty:
        print("Top stocks:")
        print(top_stocks[['code', 'name', 'total_score', 'technical_score', 'fundamental_score']].head())
    else:
        print("No stocks found matching criteria")
